from typing import List
from FlagEmbedding import FlagAutoModel
from langchain.embeddings.base import Embeddings
import logging


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class FlagEmbeddingsWrapper(Embeddings):
    def __init__(self):
        
        self.model = None
        self._initialize_model()

    def _initialize_model(self, retry_count=3):
        model_name = 'BAAI/bge-small-en-v1.5'
        fallback_model = 'BAAI/bge-small-en-v1.5'
        
        for attempt in range(retry_count):
            try:
                self.model = FlagAutoModel.from_finetuned(
                    model_name,
                    query_instruction_for_retrieval="为这个句子生成表示以用于检索相关段落：",
                    use_fp16=True
                )
                logger.info(f"成功加载模型 {model_name}")
                return
            except Exception as e:
                logger.warning(f"模型加载失败（尝试 {attempt+1}/{retry_count}）: {str(e)}")
                if attempt == retry_count - 1:
                    logger.warning(f"将回退到轻量级模型 {fallback_model}")
                    self.model = FlagAutoModel.from_finetuned(fallback_model)
                    return

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        return self.model.encode(texts, batch_size=32).tolist()  # 修改2：转换为list
    
    def embed_query(self, text: str) -> List[float]:
        return self.model.encode([text], batch_size=1)[0].tolist()  # 修改3：转换为list